import cv2

from ultralytics import YOLO
from ultralytics.solutions import speed_estimation

model = YOLO("yolov8n.pt")
names = model.model.names

cap = cv2.VideoCapture(r"D:\yolo\Speed-detection-of-vehicles-main\highway.mp4")
assert cap.isOpened(), "Error reading video file"
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))

# Video writer
video_writer = cv2.VideoWriter("speed_estimation1.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))

line_pts = [(w/2,0),(w/2,h)] #这里定义了两条线，用于车道线或其他参考线
#line_pts = [(0,h/2),(w,h/2)] # 定义一条水平线

# Init speed-estimation obj
speed_obj = speed_estimation.SpeedEstimator()
speed_obj.set_args(
    reg_pts=line_pts,
    names=names,
    view_img=True,
    spdl_dist_thresh=10,
    line_thickness=2
)

while cap.isOpened():
    success, im0 = cap.read()
    if not success:
        print("Video frame is empty or video processing has been successfully completed.")
        break

    tracks = model.track(im0,persist=True, conf=0.3,iou=0.5,show=False,
                         classes=[0,1,2,3])

    im0 = speed_obj.estimate_speed(im0, tracks)
    video_writer.write(im0)

cap.release()
video_writer.release()
cv2.destroyAllWindows()
